#include "udp_handler.hpp"
#include <cstring>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>

UDPHandler::UDPHandler(const char* server_ip, int server_port)
    : server_ip_(server_ip)
    , server_port_(server_port)
    , sockfd_(-1)
    , is_initialized_(false) {
}

UDPHandler::~UDPHandler() {
    if (sockfd_ >= 0) {
        close_connection();
    }
}

bool UDPHandler::init() {
    sockfd_ = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd_ < 0) {
        std::cerr << "Error creating socket!" << std::endl;
        return false;
    }

    struct sockaddr_in server_addr;
    std::memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(server_port_);
    server_addr.sin_addr.s_addr = inet_addr(server_ip_);

    if (bind(sockfd_, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        std::cerr << "Bind failed! Error: " << strerror(errno) << std::endl;
        close(sockfd_);
        return false;
    }

    is_initialized_ = true;
    return true;
}

bool UDPHandler::receive_image(cv::Mat& image) {
    if (!is_initialized_) return false;

    char buffer[BUFFER_SIZE];
    struct sockaddr_in client_addr;
    socklen_t addr_len = sizeof(client_addr);
    
    // 清空之前的数据
    image_data_.clear();
    
    while (true) {
        ssize_t bytes_received = recvfrom(sockfd_, buffer, BUFFER_SIZE, 0, 
                                        (struct sockaddr*)&client_addr, &addr_len);
        if (bytes_received < 0) {
            std::cerr << "Error: Receiving data failed" << std::endl;
            return false;
        }

        // 检查数据大小限制
        if (image_data_.size() + bytes_received > 1024 * 1024) {  // 1MB限制
            std::cerr << "Error: Data size too large" << std::endl;
            image_data_.clear();
            return false;
        }

        // 添加新数据
        image_data_.insert(image_data_.end(), buffer, buffer + bytes_received);

        // 检查JPEG头部标记
        if (image_data_.size() >= 2 && 
            (unsigned char)image_data_[0] != 0xFF || 
            (unsigned char)image_data_[1] != 0xD8) {
            std::cerr << "Error: Invalid JPEG header" << std::endl;
            image_data_.clear();
            return false;
        }

        // 检查是否接收到完整的JPEG图像
        if (image_data_.size() >= 2 && 
            (unsigned char)image_data_[image_data_.size() - 2] == 0xFF && 
            (unsigned char)image_data_[image_data_.size() - 1] == 0xD9) {
            
            try {
                image = cv::imdecode(image_data_, cv::IMREAD_COLOR);
                if (!image.empty()) {
                    //std::cout << "Successfully decoded image: " 
                             //<< image.size() << std::endl;
                    image_data_.clear();
                    return true;
                }
            } catch (const cv::Exception& e) {
                std::cerr << "OpenCV error: " << e.what() << std::endl;
            }
            
            std::cerr << "Error: Failed to decode image" << std::endl;
            image_data_.clear();
            return false;
        }
    }
}

bool UDPHandler::send_result(const std::string& result) {
    if (!is_initialized_) return false;
    
    // 设置PC的地址信息
    struct sockaddr_in pc_addr;
    std::memset(&pc_addr, 0, sizeof(pc_addr));
    pc_addr.sin_family = AF_INET;
    pc_addr.sin_port = htons(5005);  // 使用相同的端口
    pc_addr.sin_addr.s_addr = inet_addr("192.168.0.1");  // PC的IP地址
    
    // 发送结果到PC
    ssize_t bytes_sent = sendto(sockfd_, result.c_str(), result.length(), 0,
                               (struct sockaddr*)&pc_addr, sizeof(pc_addr));
                               
    if (bytes_sent < 0) {
        std::cerr << "Failed to send result to PC" << std::endl;
        return false;
    }
    
    return true;
}

void UDPHandler::close_connection() {
    if (sockfd_ >= 0) {
        close(sockfd_);
        sockfd_ = -1;
    }
    is_initialized_ = false;
}
